import bpy
import numpy as np
import math
import json
import os

bl_info = {
    "name": "Math Curve",
    "description": "Create a custom math curve use expression",
    "author": "moshi",
    "version": (0, 0, 1),
    "blender": (2, 80, 0),  # Blender version compatibility
    "location": "View3D > Add > Mesh > Math Curve",
    "warning": "",  # Used for warning messages for outdated plugins
    "category": "Add Mesh"
}
def resample_points(points, num_samples):
    """
    对输入的点列表进行重采样，使采样点间距相等
    :param points: 原始点列表
    :param num_samples: 重采样后的点数
    :return: 重采样后的点列表
    """
    # 计算原始点之间的累计距离
    distances = [0]
    for i in range(1, len(points)):
        dist = np.linalg.norm(points[i] - points[i - 1])
        distances.append(distances[-1] + dist)
    total_distance = distances[-1]

    # 计算重采样点的目标间距
    target_spacing = total_distance / (num_samples - 1)

    resampled_points = []
    current_distance = 0
    index = 0
    for i in range(num_samples):
        target_distance = i * target_spacing
        while index < len(distances) - 1 and distances[index + 1] < target_distance:
            index += 1
        if index == len(distances) - 1:
            resampled_points.append(points[-1])
        else:
            # 线性插值计算重采样点的位置
            alpha = (target_distance - distances[index]) / (distances[index + 1] - distances[index])
            resampled_point = (1 - alpha) * points[index] + alpha * points[index + 1]
            resampled_points.append(resampled_point)

    return resampled_points


class CreateMathCurve(bpy.types.Operator):
    """Create a custom math curve"""
    # bl_idname 用于在 Blender 中注册操作 必须只能包含一个.
    bl_idname = "curve.create_math_curve"
    bl_label = "Create Math Curve"
    bl_options = {'REGISTER', 'UNDO', 'PRESET'}

    # 可调节参数
    t_start: bpy.props.FloatProperty(
        name="Time Start",
        description="Start time of the simulation",
        default=0.0
    )
    t_end: bpy.props.FloatProperty(
        name="Time End",
        description="End time of the simulation",
        default=10.0
    )
    sample_interval: bpy.props.FloatProperty(
        name="Sample Interval",
        description="Interval between sample points",
        default=0.01,
        min=0.0001
    )
    x_expr: bpy.props.StringProperty(
        name="X Expression",
        description="Expression for x in terms of t",
        default="sin(t)"
    )
    y_expr: bpy.props.StringProperty(
        name="Y Expression",
        description="Expression for y in terms of t",
        default="t"
    )
    z_expr: bpy.props.StringProperty(
        name="Z Expression",
        description="Expression for z in terms of t",
        default="0"
    )
    resample_points_count: bpy.props.IntProperty(
        name="Resample Points Count",
        description="Number of points after resampling. Set to 0 to disable resampling.",
        default=0,
        min=0
    )
    points_count  = 0
    is_closed: bpy.props.BoolProperty(
        name="Is Closed Curve",
        description="Whether the curve is closed",
        default=False
    )

    def execute(self, context):
        # 初始化空的点列表
        points = []
        t = self.t_start
        while t < self.t_end:
            try:
                # 创建一个局部命名空间，包含数学函数和当前的 t 值
                namespace = {
                    't': t,
                    **math.__dict__
                }
                # 计算 x, y, z 的值
                x = eval(self.x_expr, namespace)
                y = eval(self.y_expr, namespace)
                z = eval(self.z_expr, namespace)
                points.append(np.array([x, y, z]))
            except Exception as e:
                self.report({'ERROR'}, f"Error evaluating expression: {e}")
                return {'CANCELLED'}
            t += self.sample_interval

        # 重采样
        if self.resample_points_count > 2:
            points = resample_points(points, self.resample_points_count)

        # 创建新的曲线对象
        curve_data = bpy.data.curves.new(name="MathCurve", type='CURVE')
        curve_data.dimensions = '3D'
        polyline = curve_data.splines.new('POLY')
        polyline.points.add(len(points) - 1)

        # 设置曲线的点
        for i, point in enumerate(points):
            polyline.points[i].co = (*point, 1)

        # 设置曲线是否闭合
        polyline.use_cyclic_u = self.is_closed

        # 从曲线数据创建对象
        curve_object = bpy.data.objects.new("MathCurveObject", curve_data)
        bpy.context.collection.objects.link(curve_object)

        # 自动聚焦到新创建的曲线对象
        bpy.ops.object.select_all(action='DESELECT')
        curve_object.select_set(True)
        context.view_layer.objects.active = curve_object
        bpy.ops.view3d.view_selected()

        # 统计点数并报告
        self.points_count = len(points)

        return {'FINISHED'}

    def draw(self, context):
        layout = self.layout
        layout.prop(self, "t_start")
        layout.prop(self, "t_end")
        layout.prop(self, "sample_interval")
        layout.prop(self, "x_expr")
        layout.prop(self, "y_expr")
        layout.prop(self, "z_expr")
        layout.prop(self, "resample_points_count")
        layout.prop(self, "is_closed")

        # 计算预计点数
        layout.label(text=f"count of total points: {self.points_count}")

        # 添加保存和加载预设的按钮
        row = layout.row()
        row.operator("object.save_math_curve_preset", text="Save Preset")
        row.operator("object.load_math_curve_preset", text="Load Preset")


class SaveMathCurvePreset(bpy.types.Operator):
    """Save the current math curve settings as a preset"""
    bl_idname = "object.save_math_curve_preset"
    bl_label = "Save Math Curve Preset"

    def execute(self, context):
        operator = context.window_manager.operators[-1] if context.window_manager.operators else None
        if operator:
            preset = {
                "t_start": operator.t_start,
                "t_end": operator.t_end,
                "sample_interval": operator.sample_interval,
                "x_expr": operator.x_expr,
                "y_expr": operator.y_expr,
                "z_expr": operator.z_expr,
                "resample_points_count": operator.resample_points_count,
                "is_closed": operator.is_closed
            }
            # 获取预设文件路径
            preset_dir = os.path.join(bpy.utils.user_resource('SCRIPTS'), "presets", "math_curve")
            if not os.path.exists(preset_dir):
                os.makedirs(preset_dir)
            preset_path = os.path.join(preset_dir, "math_curve_preset.json")
            with open(preset_path, 'w') as f:
                json.dump(preset, f, indent=4)
            self.report({'INFO'}, "Preset saved successfully.")
        return {'FINISHED'}


class LoadMathCurvePreset(bpy.types.Operator):
    """Load a saved math curve preset"""
    bl_idname = "object.load_math_curve_preset"
    bl_label = "Load Math Curve Preset"

    def execute(self, context):
        preset_dir = os.path.join(bpy.utils.user_resource('SCRIPTS'), "presets", "math_curve")
        preset_path = os.path.join(preset_dir, "math_curve_preset.json")
        if os.path.exists(preset_path):
            with open(preset_path, 'r') as f:
                preset = json.load(f)
            operator = context.window_manager.operators[-1] if context.window_manager.operators else None
            if operator:
                operator.t_start = preset["t_start"]
                operator.t_end = preset["t_end"]
                operator.sample_interval = preset["sample_interval"]
                operator.x_expr = preset["x_expr"]
                operator.y_expr = preset["y_expr"]
                operator.z_expr = preset["z_expr"]
                operator.resample_points_count = preset["resample_points_count"]
                operator.is_closed = preset["is_closed"]
            self.report({'INFO'}, "Preset loaded successfully.")
        else:
            self.report({'ERROR'}, "Preset file not found.")
        return {'FINISHED'}


def menu_func(self, context):
    self.layout.operator(CreateMathCurve.bl_idname)


def register():
    bpy.utils.register_class(CreateMathCurve)
    bpy.utils.register_class(SaveMathCurvePreset)
    bpy.utils.register_class(LoadMathCurvePreset)
    bpy.types.VIEW3D_MT_curve_add.append(menu_func)


def unregister():
    bpy.utils.unregister_class(CreateMathCurve)
    bpy.utils.unregister_class(SaveMathCurvePreset)
    bpy.utils.unregister_class(LoadMathCurvePreset)
    bpy.types.VIEW3D_MT_curve_add.remove(menu_func)


if __name__ == "__main__":
    register()